import torch

from ultralytics.nn.tasks import  attempt_load_weights
from PlateOcrNetColored import PlateOcrNetColored

def load_yolov8_model(weights, device):
    """
    加载 YOLOv8 模型。

    参数：
    weights：YOLOv8 模型的权重文件路径。
    device：模型加载到的设备。

    返回值：
    加载的模型。
    """
    # 调用 attempt_load_weights 函数加载 FP32 类型的模型
    return attempt_load_weights(weights, device=device)

def init_ocr_model(device, model_path, plateName,is_color=False):
    """
    初始化模型，加载预训练权重，并将模型移动到指定设备上。

    参数：
    device：指定的设备，如 'cuda' 或 'cpu'。
    model_path：预训练模型的路径。
    is_color：是否为彩色图像（默认为 False）。

    返回值：
    初始化后的模型。
    """

    # 加载预训练模型的权重
    check_point = torch.load(model_path, map_location=device)
    model_state = check_point['state_dict']
    cfg = check_point['cfg']
    
    color_classes = 0
    if is_color:
        color_classes = 5  # 颜色类别数

    # 初始化模型
    model = PlateOcrNetColored(num_classes=len(plateName), export=True, cfg=cfg, color_num=color_classes)
   
    # 加载模型权重
    model.load_state_dict(model_state, strict=False)
    
    # 将模型移动到指定设备上
    model.to(device)
    
    # 设置模型为评估模式
    model.eval()

    return model